from fastapi import FastAPI, WebSocket, WebSocketDisconnect, BackgroundTasks
from Agents import AgentClass
from AddDoc import AddDocClass
import uuid
from langsmith.wrappers import wrap_openai
from langsmith import traceable

# 工具
import os
from dotenv import load_dotenv

load_dotenv(verbose=True)
api_key = os.getenv("DASHSCOPE_API_KEY")
api_base = os.getenv("TONGYI_API_BASE")
user_agent = os.getenv("USER_AGENT")
print(api_key)

app = FastAPI()

@app.get("/")
def read_root():
    return {"message": "Hello, World!"}


@app.post("/chat")
def chat(query: str, background_tasks: BackgroundTasks):
    agent = AgentClass()
    msg = agent.run_agent(query)
    unique_id = str(uuid.uuid4())#生成唯一标识符
    background_tasks.add_task(agent.background_voice_synthesis, msg["output"], unique_id)
    return {"msg": msg, "id": unique_id}

@app.post("/add_urls")
async def add_urls(url: str):
    add_doc = AddDocClass()
    await add_doc.add_urls(url)
    return {"ok": "添加成功"}


@app.post("/add_pdfs")
def add_pdfs():
    return {"message": "PDFs added successfully"}


@app.post("/add_texts")
def add_text():
    return {"message": "Text added successfully"}


@app.websocket("/ws")
async def websocket_endpoint(webSocket: WebSocket):
    await webSocket.accept()
    try:
        while True:
            data = await webSocket.receive_text()
            await webSocket.send_text(f"Message text was: {data}")
    except WebSocketDisconnect:
        print("WebSocket connection closed")
        await webSocket.close()


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8001)